import torch
import torch.nn as nn
import torch.nn.functional as F


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1)
        self.fc1 = nn.Linear(in_features=9216, out_features=128)
        self.fc2 = nn.Linear(in_features=128, out_features=10)

    def forward(self, x):
        # input.shape = [N, 3, 28, 28]
        x = self.conv1(x)  # [N, 32, 26, 26]
        x = F.relu(x)
        x = self.conv2(x)  # [N, 64, 24, 24]
        x = F.relu(x)
        x = F.max_pool2d(x, 2)  # [N, 64, 12, 12]
        x = torch.flatten(input=x, end_dim=1)  # [N, 9216]
        x = self.fc1(x)  # [N, 128]
        x = F.relu(x)
        x = self.fc2(x)  # [N, 10]
        output = F.log_softmax(input=x, dim=1)
        return output


def train(train_loader, model, criterion, optimizer):
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        outputs = model(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
# TODO:添加模型验证以及主函数
